import os
import numpy as np
import pyexr
from pathlib import Path
from tqdm import tqdm
import glob

def analyze_depth_dataset():
    """分析整个数据集的深度范围"""
    dataset_path = '/baai-cwm-vepfs/public_data/scenes/trans_licheng/glassverse_v0_120_views_hdri'
    path = Path(dataset_path)
    
    print("=== 深度数据集分析 ===")
    print(f"数据集路径: {path.absolute()}")
    
    # 找到所有深度文件
    depth_files = list(path.rglob('*_depth.exr'))
    print(f"找到 {len(depth_files)} 个深度文件")
    
    if len(depth_files) == 0:
        print("未找到深度文件")
        return
    
    # 统计深度范围
    min_depth = float('inf')
    max_depth = float('-inf')
    depth_values = []
    valid_files = 0
    
    print("正在分析深度文件...")
    
    # 分析所有深度文件
    # sample_files = depth_files

    ratio = 0.02

    sample_length = int(len(depth_files) * ratio)
    sample_files = np.random.choice(depth_files, sample_length, replace=False)
    print(f"采样了 {sample_length} 个深度文件")


    
    for depth_file in tqdm(sample_files, desc="分析深度文件"):
        try:
            # 加载深度文件
            raw_depth = pyexr.open(str(depth_file)).get()
            
            # 去掉最后一个通道（alpha/mask），取平均值
            if raw_depth.shape[-1] > 1:
                raw_depth = raw_depth[:,:,:-1].mean(-1)
            
            # 过滤掉无效值（NaN, inf, 负数等）
            valid_mask = np.isfinite(raw_depth) & (raw_depth > 0)
            if np.any(valid_mask):
                valid_depth = raw_depth[valid_mask]
                
                file_min = np.min(valid_depth)
                file_max = np.max(valid_depth)
                if file_max > 1000:
                    print(depth_file)
                
                min_depth = min(min_depth, file_min)
                max_depth = max(max_depth, file_max)
                
                depth_values.extend(valid_depth.flatten())
                valid_files += 1
                
        except Exception as e:
            print(f"处理文件 {depth_file} 时出错: {e}")
            continue
    
    if valid_files == 0:
        print("没有找到有效的深度文件")
        return
    
    # 计算统计信息
    depth_values = np.array(depth_values)
    
    print(f"\n📊 深度统计信息:")
    print(f"  分析的文件数: {valid_files}")
    print(f"  总深度值数: {len(depth_values):,}")
    print(f"  最小深度: {min_depth:.6f}")
    print(f"  最大深度: {max_depth:.6f}")
    print(f"  深度范围: {max_depth - min_depth:.6f}")
    print(f"  平均深度: {np.mean(depth_values):.6f}")
    print(f"  深度中位数: {np.median(depth_values):.6f}")
    print(f"  深度标准差: {np.std(depth_values):.6f}")
    
    # 计算分位数
    percentiles = [1, 5, 10, 25, 50, 75, 90, 95, 99]
    print(f"\n📈 深度分位数:")
    for p in percentiles:
        value = np.percentile(depth_values, p)
        print(f"  {p}%: {value:.6f}")
    
    # 保存统计结果
    stats = {
        'min_depth': float(min_depth),
        'max_depth': float(max_depth),
        'depth_range': float(max_depth - min_depth),
        'mean_depth': float(np.mean(depth_values)),
        'median_depth': float(np.median(depth_values)),
        'std_depth': float(np.std(depth_values)),
        'total_values': len(depth_values),
        'valid_files': valid_files,
        'percentiles': {str(p): float(np.percentile(depth_values, p)) for p in percentiles}
    }
    
    import json
    with open('depth_statistics.json', 'w') as f:
        json.dump(stats, f, indent=2)
    
    print(f"\n💾 统计结果已保存到 depth_statistics.json")
    
    return stats

def analyze_single_scene_depth(scene_path):
    """分析单个场景的深度数据"""
    path = Path(scene_path)
    
    if not path.exists():
        print(f"场景路径不存在: {scene_path}")
        return
    
    print(f"=== {path.name} 场景深度分析 ===")
    
    # 找到该场景的所有深度文件
    depth_files = list(path.rglob('*_depth.exr'))
    print(f"找到 {len(depth_files)} 个深度文件")
    
    if len(depth_files) == 0:
        print("未找到深度文件")
        return
    
    scene_min_depth = float('inf')
    scene_max_depth = float('-inf')
    scene_depth_values = []
    
    for depth_file in depth_files:
        try:
            # 加载深度文件
            raw_depth = pyexr.open(str(depth_file)).get()
            
            # 去掉最后一个通道（alpha/mask），取平均值
            if raw_depth.shape[-1] > 1:
                raw_depth = raw_depth[:,:,:-1].mean(-1)
            
            # 过滤掉无效值
            valid_mask = np.isfinite(raw_depth) & (raw_depth > 0)
            if np.any(valid_mask):
                valid_depth = raw_depth[valid_mask]
                
                file_min = np.min(valid_depth)
                file_max = np.max(valid_depth)
                
                scene_min_depth = min(scene_min_depth, file_min)
                scene_max_depth = max(scene_max_depth, file_max)
                
                scene_depth_values.extend(valid_depth.flatten())
                
        except Exception as e:
            print(f"处理文件 {depth_file} 时出错: {e}")
            continue
    
    if len(scene_depth_values) == 0:
        print("没有找到有效的深度数据")
        return
    
    scene_depth_values = np.array(scene_depth_values)
    
    print(f"\n📊 {path.name} 场景深度统计:")
    print(f"  深度文件数: {len(depth_files)}")
    print(f"  有效深度值数: {len(scene_depth_values):,}")
    print(f"  最小深度: {scene_min_depth:.6f}")
    print(f"  最大深度: {scene_max_depth:.6f}")
    print(f"  深度范围: {scene_max_depth - scene_min_depth:.6f}")
    print(f"  平均深度: {np.mean(scene_depth_values):.6f}")
    print(f"  深度中位数: {np.median(scene_depth_values):.6f}")
    print(f"  深度标准差: {np.std(scene_depth_values):.6f}")

if __name__ == "__main__":
    # 分析整个数据集
    dataset_stats = analyze_depth_dataset()
    
    # 分析单个场景示例
    demo_scene = '/baai-cwm-vepfs/public_data/scenes/trans_licheng/glassverse_v0_120_views_hdri/000000'
    analyze_single_scene_depth(demo_scene) 